# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#     http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Code is heavily based on paper "Geometry-Informed Neural Operator for Large-Scale 3D PDEs", we use paddle to reproduce the results of the paper


import paddle
from functools import partial
from .base_model import BaseModel
from paddle import nn


class ResBlock(paddle.nn.Layer):
    def __init__(self, channel, **kwargs):
        super(ResBlock, self).__init__(**kwargs)
        self.channel = channel
        self._conv_1 = nn.Conv3D(self.channel, self.channel, kernel_size=3, padding='same')
        self._conv_2 = nn.Conv3D(self.channel, self.channel, kernel_size=3, padding='same')
        self._gelu = nn.GELU()

    def forward(self, x):
        x_ini = x
        x = self._conv_1(x)
        x = self._gelu(x)
        x = self._conv_2(x)
        return x+x_ini
class Res3D(BaseModel):

    def __init__(self, in_channels,hidden_channels,num_blocks=12):
        super(Res3D, self).__init__()
        self.conv_in=nn.Conv3D(in_channels, hidden_channels, kernel_size=3, padding='same')
        self.blocks = nn.Sequential()
        for block_id in range(num_blocks):
            self.blocks.add_sublayer('block_{}'.format(block_id), ResBlock(channel=hidden_channels))
    def forward(self, inputs):
        x = self.conv_in(inputs)
        x = self.blocks(x)
        return x
class UNet3DWithRes3D(Res3D):

    def __init__(self, in_channels: int, out_channels: int, hidden_channels:
        int, num_levels: int, use_position_input: bool=True):
        super(UNet3DWithRes3D, self).__init__(in_channels,
            hidden_channels, num_blocks=num_levels)
        self.final_mlp = paddle.nn.Sequential(paddle.nn.Linear(in_features=
            hidden_channels, out_features=hidden_channels), paddle.nn.GELU(
            ), paddle.nn.Linear(in_features=hidden_channels, out_features=
            out_channels))
        self.use_position_input = use_position_input

    def forward(self, x, output_points):
        x = super(UNet3DWithRes3D, self).forward(x)
        output_points = output_points.unsqueeze(axis=2).unsqueeze(axis=2)
        x = paddle.nn.functional.grid_sample(x=x, grid=output_points,
            align_corners=False)
        x = x.squeeze(axis=3).squeeze(axis=3)
        x = x.transpose(perm=[0, 2, 1])
        x = self.final_mlp(x)
        return x

    def data_dict_to_input(self, data_dict):
        input_grid_features = data_dict['sdf'].unsqueeze(axis=1)

        if self.use_position_input:
            grid_points = data_dict['sdf_query_points']
            input_grid_features = paddle.concat(x=(input_grid_features,
                grid_points), axis=1)
        output_points = data_dict['vertices']
        return input_grid_features, output_points

    @paddle.no_grad()
    def eval_dict(self, data_dict, loss_fn=None, decode_fn=None, **kwargs):
        input_grid_features, output_points = self.data_dict_to_input(data_dict)
        p_predict = self(input_grid_features, output_points)
        p_predict_decode = decode_fn(p_predict)
        out_dict = {}
        out_dict['pressure'] = p_predict
        out_dict['p_predict_decode'] = p_predict_decode

        if loss_fn is not None:
            p_label = data_dict['pressure'].unsqueeze(-1)
            p_label_decode = decode_fn(p_label)
            out_dict['l2'] = loss_fn(p_predict, p_label)
            out_dict['l2_decoded'] = loss_fn(p_predict_decode, p_label_decode)
            out_dict['p_label_decode'] = p_label_decode
        return out_dict

    def loss_dict(self, data_dict, loss_fn=None, **kwargs):
        input_grid_features, output_points = self.data_dict_to_input(data_dict)

        vert_sdf = self(input_grid_features, output_points)
        gt_pressure = data_dict['pressure'].unsqueeze(-1)

        if loss_fn is None:
            loss_fn = self.loss
        return {'loss': loss_fn(vert_sdf, gt_pressure)}

